from pyVim.connect import SmartConnect, Disconnect
from pyVmomi import vim
import ssl
import re
import math
import time

"""
Get VM Data From vCenter
"""

def connect_to_vcenter(host, user, pwd, port):
    """
    paras:
        host: vCenter IP or FQDN
        user: vCenter username
        pwd: vCenter password
        port: vCenter port
    """
    # 创建SSL上下文，禁用证书验证
    context = ssl._create_unverified_context()
    # 连接到vCenter
    try:
        si = SmartConnect(host=host, user=user, pwd=pwd, port=port, sslContext=context)
        assert si  
              
        return si
    
    except Exception as e:
        return "vCenter连接失败:" + str(e)

def disconnect_from_vcenter(si):
    try:
        # 断开连接
        Disconnect(si)
        return "断开连接成功"
        
    except Exception as e:
        return "断开连接失败: " + str(e)

def split_hostname_domain(fullname):
    """
    拆分主机名和域名,若是IP地址则直接保留原样返回
    """
    # 判断是IP还是域名
    if '.' in fullname and fullname.count('.') == 3:
        return {'hostname': fullname, 'domain': None}
    # 分割字符串为两部分：hostname和domain
    else:
        parts = fullname.split('.', 1)
        hostname = parts[0]  # 最后一个点号之前的所有部分
        domain = parts[1] if len(parts) > 1 else None  # 最后一个点号之后的部分，如果不存在则为空字符串
        return {'hostname': hostname, 'domain': domain}

def split_datastore(filename):
    """
    处理datastore,截取vSan的名字
    """
    match = re.match(r'\[(.*?)\]', filename)
    if match:
        content = match.group(1)
        if content.startswith("VxRail-Virtual-SAN-Datastore"):
            content = "VxRail-Virtual-SAN-Datastore"
            return content
        else:
            return content
    else:
        return None
    
def get_vm_ip(vm_networks):
    """
    获取VM的各类IP地址
    """
    ipaddress = None
    backup_ip = None
    sub_ipaddress = None
    ip_list = []

    for net in vm_networks:
        if net.ipConfig is not None and net.ipConfig.ipAddress is not None:
            for ip in net.ipConfig.ipAddress:
                if ip.prefixLength in (24, 28):
                    if ip.ipAddress.startswith('175.'):
                        backup_ip = ip.ipAddress
                    else:
                        ip_list.append(ip.ipAddress)
    
    if ip_list != []:
        if len(ip_list) == 1:
            ipaddress = ip_list[0]
        else:
            ip_list.sort()
            ipaddress = ip_list[0]
            sub_ip = ip_list[1:]
            sub_ipaddress = ','.join(sub_ip)

    return ipaddress, backup_ip, sub_ipaddress
    
def get_vm_info(si, log_path):
    """
    获取需要的VM信息
    """
    # 定义一个空列表来存储VM信息
    vms_data = []
    try:
        # 获取全部虚拟机
        content = si.RetrieveContent()
        container = content.rootFolder
        viewType = [vim.VirtualMachine]
        recursive = True
        containerView = content.viewManager.CreateContainerView(container, viewType, recursive)
        # 遍历虚拟机
        virtual_machines = containerView.view

        for vm in virtual_machines:
            vm_power_state = vm.summary.runtime.powerState
            if vm_power_state == 'poweredOn' and 'vCLS' not in vm.name:
                vmid = vm._moId
                name = vm.name
                nics = sum(1 for device in vm.config.hardware.device if isinstance(device, vim.vm.device.VirtualEthernetCard))
                vm_nets = get_vm_ip(vm.guest.net) if vm.guest.net is not None else (None, None, None)
                ipaddress, backup_ip, sub_ipaddress = vm_nets
                vmtool= vm.guest.toolsStatus
                hostname_full = vm.summary.guest.hostName
                hostname = split_hostname_domain(hostname_full)['hostname'] if hostname_full else ''
                domain = split_hostname_domain(hostname_full)['domain'] if hostname_full else ''
                os_version = vm.config.guestFullName
                if 'Windows' in os_version:
                    os_type = 'Windows'
                else:
                    os_type = 'Linux'
                cpu_cores = vm.summary.config.numCpu
                mem_size = math.ceil(vm.summary.config.memorySizeMB / 1024)
                disk_num = sum(1 for device in vm.config.hardware.device if isinstance(device, vim.vm.device.VirtualDisk))
                disk_size = math.ceil(sum([device.capacityInKB for device in vm.config.hardware.device if isinstance(device, vim.vm.device.VirtualDisk)]) / (1024 * 1024))
                vmhost_full = vm.runtime.host.name
                vmhost = split_hostname_domain(vmhost_full)['hostname'] if vmhost_full else ''
                vmdatastore = None
                for device in vm.config.hardware.device:
                    if isinstance(device, vim.vm.device.VirtualDisk):
                        backing = device.backing
                        if isinstance(backing, vim.vm.device.VirtualDisk.FlatVer2BackingInfo):
                            filename =  backing.fileName
                            vmdatastore = split_datastore(filename)
                            break
                        
                vms_data.append({
                    "vmid": vmid, "name": name, "ipaddress": ipaddress, "backup_ip": backup_ip, "sub_ipaddress": sub_ipaddress, "nics": nics, "vmtool": vmtool, "hostname": hostname, "domain": domain, "os_type": os_type,
                    "os_version": os_version, "cpu_cores": cpu_cores, "mem_size": mem_size, "disk_num": disk_num, "disk_size": disk_size, "vmhost": vmhost, "vmdatastore": vmdatastore,                    
                })

        containerView.Destroy()
        message = "Get vm info total: %s" %len(vms_data)
        return vms_data
    except Exception as e:
        message = "Fail to get vm info:" + str(e)
    with open(log_path, 'ab') as f:
        log = 'Time:%s \t Result:%s \n' % (time.strftime('%Y-%m-%d %H:%M:%S'), message)
        f.write(log.encode())
